﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Win32.SafeHandles;

namespace JasLib.Threading
{
    public static class ThreadingExtensions
    {
        /**********************************************************************************/
        public static Task ForEachAsync<T>(
            this IEnumerable<T> source,
            Func<T, Task> asyncFunc,
            CancellationToken cancellationToken = default(CancellationToken))
        {
            List<Task> tasks = new List<Task>();

            /// TODO: This can't be right. Fix it.
            foreach (T t in source)
            {
                Task task = Task.Factory.StartNew(obj => asyncFunc((T)obj), t, cancellationToken);
                tasks.Add(task);
            }

            if (tasks.Count > 1)
            {
                /// WhenAll() internally converts the enumerable into an array using the ToArray extension method.
                /// It's more efficient for us to do it up front using the List<T>.ToArray instance method.
                return TaskEx.WhenAll(tasks.ToArray());
            }
            else if (tasks.Count == 1)
                return tasks[0];
            else
                return ThreadingUtilities.CompletedTask;
        }

        /**********************************************************************************/
        public static TaskAwaiter<bool> GetAwaiter(this SafeWaitHandle waitHandle)
        {
            return WaitOneAsync(waitHandle).GetAwaiter();
        }

        /**********************************************************************************/
        public static TaskAwaiter<bool> GetAwaiter(this WaitHandle waitHandle)
        {
            return WaitOneAsync(waitHandle).GetAwaiter();
        }

        /**********************************************************************************/
        public static TaskAwaiter GetAwaiter(this TimeSpan timespan)
        {
            return ThreadingUtilities.SleepAsync(timespan).GetAwaiter();
        }

        /**********************************************************************************/
        public static TaskAwaiter GetAwaiter(this DateTime date)
        {
            return ThreadingUtilities.SleepUntilAsync(date).GetAwaiter();
        }

        /**********************************************************************************/
        public static Awaitable SwitchToAsync(this ISynchronizeInvoke invoker)
        {
            return AsyncThreadSwitcher.SwitchToSynchronizeInvokeAsync(invoker);
        }

        /**********************************************************************************/
        public static Awaitable SwitchToAsync(this SynchronizationContext context)
        {
            return AsyncThreadSwitcher.SwitchToSynchronizationContextAsync(context);
        }

        /***************************************************************************/
        /// <summary>
        /// WaitHandle is strangely abstract, so we only needed to inherit to use it.
        /// </summary>
        private class NonAbstractWaitHandle : WaitHandle
        {
            public NonAbstractWaitHandle(SafeWaitHandle safeWaitHandle)
            {
                this.SafeWaitHandle = safeWaitHandle;
                return;
            }
        }

        public static Task<bool> WaitOneAsync(this SafeWaitHandle safeWaitHandle)
        {
            NonAbstractWaitHandle waitHandle = new NonAbstractWaitHandle(safeWaitHandle);
            return WaitOneAsync(waitHandle);
        }

        /**********************************************************************************/
        private class WaitOneAsyncState
        {
            private WaitHandle _waitHandle;
            private CancellationTokenRegistration _cancellationTokenRegistration;
            private RegisteredWaitHandle _registeredWaitHandle = null;
            private TaskCompletionSource<bool> _tcs = new TaskCompletionSource<bool>();

            public Task<bool> Task
            {
                get { return _tcs.Task; }
            }

            public WaitOneAsyncState(
                WaitHandle waitHandle,
                int iMillisecondsTimeoutInterval,
                CancellationToken cancellationToken)
            {
                _waitHandle = waitHandle;

                if (cancellationToken.CanBeCanceled)
                    _cancellationTokenRegistration = cancellationToken.Register(this.OnCancelled);

                _registeredWaitHandle = ThreadPool.RegisterWaitForSingleObject(
                    _waitHandle, this.OnSignalled, null, iMillisecondsTimeoutInterval, true);

                return;
            }

            private void OnCancelled()
            {
                _tcs.TrySetException(new TaskCanceledException("WaitOneAsync aborted due to cancellation."));
                DetachAll();
                return;
            }

            private void OnSignalled(object objStateParam, bool timedOut)
            {
                _tcs.TrySetResult(timedOut);
                DetachAll();
                return;
            }

            public void DetachAll()
            {
                /// If CancellationTokenRegistration were a reference type,
                /// then we could do Interlocked.Exchange() on it too,
                /// and then two threads that hit this method at the exact same time might neatly diverge
                /// to where one frees the registered wait handle and the other frees the cancellation registration.
                /// But because this is not the case, we use the registered wait handle exchange
                /// as the "lock" for the whole thing.
                var registeredWaitHandle = Interlocked.Exchange(ref _registeredWaitHandle, null);
                if (registeredWaitHandle != null)
                {
                    registeredWaitHandle.Unregister(_waitHandle);
                    _cancellationTokenRegistration.Dispose();
                }
                return;
            }
        }

        public static Task<bool> WaitOneAsync(this WaitHandle waitHandle, int iMillisecondsTimeoutInterval = -1, CancellationToken cancellationToken = default(CancellationToken))
        {
            /// Ideally we could even prevent the post-to-SynchronizationContext if either of these conditions were signalled.

            cancellationToken.ThrowIfCancellationRequested();

            if (iMillisecondsTimeoutInterval == 0)
            {
                var tcs = new TaskCompletionSource<bool>();
                try
                {
                    tcs.TrySetResult(waitHandle.WaitOne(0));
                }
                catch (Exception ex)
                {
                    tcs.TrySetException(ex);
                }
                return tcs.Task;
            }

            return new WaitOneAsyncState(waitHandle, iMillisecondsTimeoutInterval, cancellationToken).Task;
        }

        public static Task<bool> WaitOneAsync(this WaitHandle waitHandle, TimeSpan waitTime, CancellationToken cancellationToken = default(CancellationToken))
        {
            return WaitOneAsync(waitHandle, (int)waitTime.TotalSeconds, cancellationToken);
        }

        /**********************************************************************************/
        public static bool WaitOneTrapAbandonment(this Mutex mutex, int millisecondsTimeout)
        {
            bool bAcquired = false;
            try
            {
                bAcquired = mutex.WaitOne(millisecondsTimeout);
            }
            catch (AbandonedMutexException)
            {
                /// This exception should not be confused to mean that the acquisition failed.
                bAcquired = true;
            }
            return bAcquired;
        }

        /**********************************************************************************/
        public static Task WhenAll(this IEnumerable<Task> tasks)
        {
            return TaskEx.WhenAll(tasks);
        }

        /**********************************************************************************/
        public static Task WhenAny(this IEnumerable<Task> tasks)
        {
            return TaskEx.WhenAny(tasks);
        }
    }
}
